# !/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName : zsj.py
# @Time     : 2024/5/17 22:34
# @Author   : Robot-Zsj
"""
description:

首先，考虑该问题的简化版：给定n个非递减的非负数序列a0，a1,。。。，an-1，找出第k个最小的子序列和。

当k=1，空序列即为答案
当k>1，令（t,i）表示以ai为最后一个元素且和为t的子序列，则问题转化为求解第k个最小的(t,i)。
我们使用一个小根堆来保存(t,i)，初始时堆中只有一个元素(a0,0)。为了能按从小到大的顺序依次获得子序列，
当我们从小根堆取出堆顶元素(t,i)时，我们需要进行一下操作

- 将ai+1拼接到子序列(t,i)后，得到新的子序列(t+ai+1,i+1)，并将它加入堆中。
- 将子序列(t,i)中的ai替换成ai+1，得到新的子序列(t+ai+1-ai,i+1)，并将它加入堆中。

那么第k-1次取出的堆顶元素对应第k个最小的子序列

根据以上讨论，我们可以求解出非递减的非负数序列的第k个最小的子序列和，而原问题给出的条件中，允许序列中有负数。
记原序列的非负数和与负数和分别为total与total-neg，我们先将原序列中的负数替换成它的绝对值，然后从小到大进行排序，最后求解第k个最小的子序列和tk。

tk对应一个子序列sk，我们将tk加上total-neg，那么tk+total-neg也对应原序列的一个子序列s'k，该子序列s'k的正数元素与sk相同，负数部分由绝对值不存在与sk的负数元素组成。
因为tk+total-neg是非递减的，所以s'k是原序列的第k个最小的子序列，tk+totalneg是原序列的第k个最小的子序列和。

同时，如果第k个最小的子序列和为tk+total-neg，并且所有序列元素和为total+total-neg，那么total+total-neg-（tk+total-neg)=total-tk即为第k个最大的子序列和。
"""
from heapq import heappop, heappush

nums = [2, 4, -2]
# 6 4 4 2 2 0 0 -2
k = 5
output = 2


def ksum(nums, k):
    n = len(nums)
    total = 0
    for i in range(n):
        if nums[i] >= 0:
            total += nums[i]
        else:
            nums[i] = -nums[i]

    nums.sort()

    ret = 0

    pq = [(nums[0], 0)]
    for j in range(2, k + 1):
        t, i = heappop(pq)
        ret = t
        if i == n - 1:
            continue
        heappush(pq, (t + nums[i + 1], i + 1))
        heappush(pq, (t - nums[i] + nums[i + 1], i + 1))
    return total - ret


res = ksum(nums, k)
print(res)
assert res == output, "答案错误"
